import os
from collections import Counter, defaultdict
import numpy as np
import matplotlib.pyplot as plt

def tmp():
    rate = 0.7

    fft_dir = './res/fft1/'
    acts = os.listdir(fft_dir)
    plt.figure(figsize=(10,14))
    for i in range(6):
        actions = defaultdict(list)
        rate = 0.7 + 0.05 * i
        # rate = 0.3 - 0.05 * i
        print(rate)
        for act in acts:
            act_dir = os.path.join(fft_dir, act)
            vids = os.listdir(act_dir)
            for vid in vids:
                vid_file = os.path.join(act_dir, vid)
                data = np.load(vid_file)
                num = 1
                for elem in data:
                    tmp = elem[:(len(elem)+1)//2]
                    if np.std(tmp) < 1:
                        continue
                        # plt.plot(range(len(tmp)), tmp)
                        # plt.savefig('hello.jpg')
                        # hello
                    total = sum(tmp)
                    thres = total * rate
                    count = 0
                    for idx in range(len(tmp)):
                        count += tmp[idx]
                        if count >= thres:
                            if idx >= num:
                                num = idx
                            break  
                # for elem in data:
                #     tmp = elem[:(len(elem)+1)//2]
                #     if np.std(tmp) < 0.1:
                #         break
                #         # plt.plot(range(len(tmp)), tmp)
                #         # plt.savefig('hello.jpg')
                #         # hello
                #     maxi = tmp[0]
                #     thres = maxi * rate
                #     for idx in range(len(tmp)):
                #         if tmp[idx] <= thres:
                #             if idx > num:
                #                 num = idx
                #             break  
                actions[act].append(len(elem) / num)


        values = []
        for act in actions:
            nums = actions[act]
            for num in nums:
                values.append(int(num+0.5))
        result = Counter(values)
        print(result)
        plt.subplot(3,2,i+1)
        plt.title('threshold {:.2f}'.format(rate))
        plt.xlabel('Sample Number')
        plt.ylabel('Patient Number')
        plt.bar(result.keys(), result.values()) 
    plt.savefig('./tmp1.jpg', bbox_inches='tight')

def integrate():
    rate = 0.7
    fft_dir = './res/fft1/'
    acts = os.listdir(fft_dir)
    plt.figure(figsize=(10,14))
    for i in range(6):
        actions = defaultdict(list)
        rate = 0.7 + 0.05 * i
        print(rate)
        for act in acts:
            act_dir = os.path.join(fft_dir, act)
            vids = os.listdir(act_dir)
            for vid in vids:
                vid_file = os.path.join(act_dir, vid)
                data = np.load(vid_file)
                num = 1
                for elem in data:
                    tmp = elem[:(len(elem)+1)//2]
                    if np.std(tmp) < 1:
                        continue
                        # plt.plot(range(len(tmp)), tmp)
                        # plt.savefig('hello.jpg')
                        # hello
                    total = sum(tmp)
                    thres = total * rate
                    count = 0
                    for idx in range(len(tmp)):
                        count += tmp[idx]
                        if count >= thres:
                            if idx >= num:
                                num = idx
                            break  
                actions[act].append(len(elem) / num)
        values = []
        for act in actions:
            nums = actions[act]
            for num in nums:
                values.append(int(num+0.5))
        result = Counter(values)
        print(result)
        plt.subplot(3,2,i+1)
        plt.title('threshold {:.2f}'.format(rate))
        plt.xlabel('Sample Number')
        plt.ylabel('Patient Number')
        plt.bar(result.keys(), result.values()) 
    plt.savefig('./integrate_rate.jpg', bbox_inches='tight')

def maxi():
    rate = 0.7
    fft_dir = './res/fft1/'
    acts = os.listdir(fft_dir)
    plt.figure(figsize=(10,14))
    for i in range(6):
        actions = defaultdict(list)
        rate = 0.3 - 0.05 * i
        print(rate)
        for act in acts:
            act_dir = os.path.join(fft_dir, act)
            vids = os.listdir(act_dir)
            for vid in vids:
                vid_file = os.path.join(act_dir, vid)
                data = np.load(vid_file)
                num = 1
                for elem in data:
                    tmp = elem[:(len(elem)+1)//2]
                    if np.std(tmp) < 1:
                        continue
                        # plt.plot(range(len(tmp)), tmp)
                        # plt.savefig('hello.jpg')
                        # hello
                    maxi = tmp[0]
                    thres = maxi * rate
                    for idx in range(len(tmp)):
                        if tmp[idx] <= thres:
                            if idx > num:
                                num = idx
                            break  
                actions[act].append(len(elem) / num)

        values = []
        for act in actions:
            nums = actions[act]
            for num in nums:
                values.append(int(num+0.5))
        result = Counter(values)
        print(result)
        plt.subplot(3,2,i+1)
        plt.title('threshold {:.2f}'.format(rate))
        plt.xlabel('Sample Number')
        plt.ylabel('Patient Number')
        plt.bar(result.keys(), result.values()) 
    plt.savefig('./maxi_rate.jpg', bbox_inches='tight')

# 只看特定动作的fft结果
def integrate_v2():
    rate = 0.7

    fft_dir = './res/fft1/'
    acts = os.listdir(fft_dir)
    inlist_acts = ['chew', 'brush_hair']
    for act in acts:
        if act not in inlist_acts:
            continue
        plt.figure(figsize=(15,10))
        act_dir = os.path.join(fft_dir, act)
        vids = os.listdir(act_dir)
        for i in range(6):
            res = []
            rate = 0.7 + 0.05 * i
            print(rate)
            for vid in vids:
                vid_file = os.path.join(act_dir, vid)
                data = np.load(vid_file)
                num = 1
                for elem in data:
                    tmp = elem[:(len(elem)+1)//2]
                    if np.std(tmp) < 1:
                        continue
                        # plt.plot(range(len(tmp)), tmp)
                        # plt.savefig('hello.jpg')
                        # hello
                    total = sum(tmp)
                    thres = total * rate
                    count = 0
                    for idx in range(len(tmp)):
                        count += tmp[idx]
                        if count >= thres:
                            if idx >= num:
                                num = idx
                            break  
                if num > 2:
                    res.append(len(elem) / num)

            values = []
            for num in res:
                values.append(int(num+0.5))
            result = Counter(values)
            print(result)
            plt.subplot(2,3,i+1)
            plt.xlim(1, 8)
            plt.ylim(0, 70)
            plt.title('threshold {:.2f}'.format(rate), fontsize=15)
            plt.xlabel('Stride', fontsize=15)
            plt.ylabel('Number of Videos', fontsize=15)
            plt.bar(result.keys(), result.values()) 
        plt.suptitle(act.replace('_', ' '), fontsize=20)
        plt.savefig('./res/rate1/{}.jpg'.format(act), bbox_inches='tight')
        print(act)


# 查看整个数据集的fft结果
def integrate_all(dataset):
    if dataset == 'HMDB51':
        fft_dir='./res/fft1/'
    else:
        fft_dir='./res/fft/'
    rate = 0.7
    acts = os.listdir(fft_dir)
    res = []
    for act in acts:
        plt.figure(figsize=(10,14))
        act_dir = os.path.join(fft_dir, act)
        vids = os.listdir(act_dir)
        for i in range(1):
            for vid in vids:
                vid_file = os.path.join(act_dir, vid)
                data = np.load(vid_file)
                num = 1
                for elem in data:
                    tmp = elem[:(len(elem)+1)//2]
                    if np.std(tmp) < 1:
                        continue
                        # plt.plot(range(len(tmp)), tmp)
                        # plt.savefig('hello.jpg')
                        # hello
                    total = sum(tmp)
                    thres = total * rate
                    count = 0
                    for idx in range(len(tmp)):
                        count += tmp[idx]
                        if count >= thres:
                            if idx >= num:
                                num = idx
                            break  
                if num > 2:
                    res.append(len(elem) / num)

    values = []
    for num in res:
        values.append(int(num+0.5))
    result = Counter(values)
    # plt.title(f'Sampling Rate of {dataset} When Threshold is {rate:.2f}')
    plt.title(f'{dataset}', fontsize=25)
    plt.xlabel('Stride', fontsize=25)
    plt.ylabel('Number of Videos', fontsize=25)
    plt.xticks(fontsize=20)
    plt.yticks(fontsize=20)
    plt.bar(result.keys(), result.values()) 
    plt.savefig(f'{dataset}.jpg', bbox_inches='tight')
    plt.close()
    print(act)


if __name__ == '__main__':
    # tmp()
    # integrate()
    # maxi()
    # integrate_v2()

    integrate_all(dataset='HMDB51')
    integrate_all(dataset='UCF101')